import os
import logging
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.vectorstores import Chroma

# 配置日志记录
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

model_embedding = "G:/M3E_model/xrunda/m3e-base"
model_kwargs = {'device': 'cuda'}

hf = HuggingFaceBgeEmbeddings(
    encode_kwargs={'normalize_embeddings': True},
    model_name=model_embedding,
    model_kwargs=model_kwargs,
    query_instruction="为这个句子生成表示以用于检索相关文章："
)

def load_vector_db(persist_directory):
    """加载已有的向量数据库。如果数据库不存在，则抛出错误。"""

    if not os.path.exists(persist_directory):
        logger.error(f"Database directory not found: {persist_directory}")
        raise FileNotFoundError(f"Database directory not found: {persist_directory}")

    try:
        # 使用 Chroma 构造函数加载数据库
        logger.info(f"Loading vector database from: {persist_directory}")
        db = Chroma(persist_directory=persist_directory, embedding_function=hf)
        logger.info("Vector database loaded successfully.")
        return db
    except Exception as e:
        # 捕获并记录任何加载错误
        logger.error(f"Error loading vector database: {e}")
        raise e  # 将异常抛出，以便上层可以处理
